from __future__ import absolute_import

import sys

from torch import nn
from torch.nn import functional as F

from config import get_args
from .stn_head import STNHead
from .tps_spatial_transformer import TPSSpatialTransformer
from .trans_str_realformer import Cnn_Trans_encoder, Trans_decoder
from ..loss.sequenceCrossEntropyLoss import (LabelSmoothCrossEntropyLoss,
                                             SequenceCrossEntropyLoss)

global_args = get_args(sys.argv[1:])


class ModelBuilder(nn.Module):
    """
    This is the integrated model.
    """

    def __init__(self, arch, rec_num_classes, sDim, attDim, max_len_labels, eos, STN_ON=False, label_smooth=False):  # label_smooth有待优化
        super(ModelBuilder, self).__init__()

        self.arch = arch
        self.rec_num_classes = rec_num_classes
        self.sDim = sDim
        self.attDim = attDim
        self.max_len_labels = max_len_labels
        self.eos = eos
        self.STN_ON = STN_ON
        self.tps_inputsize = global_args.tps_inputsize

        self.encoder = Cnn_Trans_encoder(in_channel=3, out_channel=512)

        self.decoder = Trans_decoder(num_class=rec_num_classes)

        # self.encoder = create(self.arch,
        #                   with_lstm=global_args.with_lstm,
        #                   n_group=global_args.n_group)
        # encoder_out_planes = self.encoder.out_planes

        # self.decoder = AttentionRecognitionHead(
        #                   num_classes=rec_num_classes,
        #                   in_planes=encoder_out_planes,
        #                   sDim=sDim,
        #                   attDim=attDim,
        #                   max_len_labels=max_len_labels)

        if label_smooth == True:
            self.rec_crit = LabelSmoothCrossEntropyLoss(0.1)
        else:
            self.rec_crit = SequenceCrossEntropyLoss()

        if self.STN_ON:
            self.tps = TPSSpatialTransformer(
                output_image_size=tuple(global_args.tps_outputsize),
                num_control_points=global_args.num_control_points,
                margins=tuple(global_args.tps_margins))
            self.stn_head = STNHead(
                in_planes=3,
                num_ctrlpoints=global_args.num_control_points,
                activation=global_args.stn_activation)

    def forward(self, input_dict):
        return_dict = {}
        return_dict['losses'] = {}
        return_dict['output'] = {}

        x, rec_targets, rec_lengths, label_mask = input_dict['images'], input_dict['rec_targets'].long(), input_dict['rec_lengths'], input_dict['label_mask']

        # rectification
        if self.STN_ON:
            # input images are downsampled before being fed into stn_head.
            stn_input = F.interpolate(x, self.tps_inputsize, mode='bilinear', align_corners=True)
            stn_img_feat, ctrl_points = self.stn_head(stn_input)
            x, _ = self.tps(x, ctrl_points)
            if not self.training:
                # save for visualization
                return_dict['output']['ctrl_points'] = ctrl_points
                return_dict['output']['rectified_images'] = x

        encoder_feats = self.encoder(x)

        encoder_feats = encoder_feats.contiguous()  # l_k,b,c

        if self.training:
            max_len = max(rec_lengths)
            rec_targets = rec_targets[:, :max_len]  # b,max_len
            rec_pred = self.decoder(encoder_feats, rec_targets, tgt_padding=label_mask)
            loss_rec = self.rec_crit(rec_pred, rec_targets, rec_lengths)
            return_dict['losses']['loss_rec'] = loss_rec
            return_dict['output']['pred_rec'] = rec_pred
        else:

            rec_pred, _, rec_score = self.decoder.inference(encoder_feats, max_len=100)
            max_len = max(rec_lengths)
            rec_targets = rec_targets[:, :max_len]  # b,max_len
            rec_logits = self.decoder(encoder_feats, rec_targets, tgt_padding=label_mask)
            # rec_logits = rec_logits[:,:max_len,:]
            loss_rec = self.rec_crit(rec_logits, rec_targets, rec_lengths)
            return_dict['losses']['loss_rec'] = loss_rec
            return_dict['output']['pred_rec'] = rec_pred
            return_dict['output']['pred_rec_score'] = rec_score

        for k, v in return_dict['losses'].items():
            return_dict['losses'][k] = v.unsqueeze(0)

        return return_dict
